
#include <cmath>
#include <jem/base/Array.h>
#include <jem/base/Error.h>
#include <jem/base/System.h>
#include <jem/base/IllegalInputException.h>
#include <jem/numeric/algebra/matmul.h>
#include <jem/numeric/algebra/MatmulChain.h>
#include <jem/util/Properties.h>
#include <jive/util/utilities.h>
#include <jive/util/XTable.h>
#include <jive/util/XDofSpace.h>
#include <jive/util/Assignable.h>
#include <jive/algebra/MatrixBuilder.h>
#include <jive/model/Model.h>
#include <jive/model/ModelFactory.h>
#include <jive/model/Actions.h>
#include <jive/model/StateVector.h>
#include <jive/geom/InternalShape.h>
#include <jive/geom/IShapeFactory.h>
#include <jive/fem/ElementGroup.h>
#include "models.h"
#include "import.h"
#include "utilities.h"
#include "Names.h"
#include "Material.h"


using jem::numeric::MatmulChain;
using jive::util::XTable;
using jive::util::XDofSpace;
using jive::util::Assignable;
using jive::algebra::MatrixBuilder;
using jive::model::Model;
using jive::geom::IShape;
using jive::fem::NodeSet;
using jive::fem::ElementSet;
using jive::fem::ElementGroup;


//=======================================================================
//   typedefs
//=======================================================================

// Some handy aliases to avoid some typing

typedef ElementSet             ElemSet;
typedef ElementGroup           ElemGroup;
typedef MatmulChain<double,2>  MChain2;
typedef MatmulChain<double,3>  MChain3;


//=======================================================================
//   class PhaseFieldModel
//=======================================================================


class PhaseFieldModel : public Model
{
 public:

  typedef PhaseFieldModel      Self;
  typedef Model             Super;

                            PhaseFieldModel

    ( const String&           name,
      const Properties&       conf,
      const Properties&       props,      // stores the input files
      const Properties&       globdat );  // stores mesh, elem groups, solution vector (later)

  virtual void              configure

    ( const Properties&       props,
      const Properties&       globdat );

  virtual void              getConfig

    ( const Properties&       conf,
      const Properties&       globdat )      const;

  virtual bool              takeAction

    ( const String&           action,
      const Properties&       params,
      const Properties&       globdat );


 protected:

  virtual                  ~PhaseFieldModel  ();


 private:

  void                      getMatrix_

    ( MatrixBuilder&          mbuilder,
      const Vector&           force,
      const Vector&           disp )         const;

  bool                      getTable_

    ( const Properties&       params,
      const Properties&       globdat );

  void                      getStress_

    ( XTable&                 table,
      const Vector&           weights,
      const Vector&           disp );


 private:

  Assignable<ElemGroup>     egroup_;
  Assignable<ElemSet>       elems_;
  Assignable<NodeSet>       nodes_;

  int                       rank_;

  Ref<IShape>               shape_;
  Ref<XDofSpace>            dofs_;
  IntVector                 dofTypes_;

  Ref<Material>             material_;
  ShapeGradsFunc            getShapeGrads_;

  double		    dummy_;
  int			    iterCount_;

};


//-----------------------------------------------------------------------
//   constructors & destructor
//-----------------------------------------------------------------------


PhaseFieldModel::PhaseFieldModel

  ( const String&      name,
    const Properties&  conf,
    const Properties&  props,
    const Properties&  globdat ) :

    Super ( name )

{
  using jem::newInstance;
  using jem::IllegalInputException;
  using jive::util::joinNames;
  using jive::geom::IShapeFactory;

  Properties    myProps = props.getProps  ( myName_ );  // extract the subset of properties contained in myName_ = model.model.solid(elements,..) from props
  Properties    myConf  = conf .makeProps ( myName_ );


  const String  context = getContext ();


  // Get the element group assigned to this model.

  egroup_ = ElemGroup::get ( myConf, myProps, globdat, context );
  // in myProps, goupset "internal" defined
  // Context contains `model.model.solid' to help find error

  elems_  = egroup_.getElements ();
  nodes_  = elems_ .getNodes    ();
  rank_   = nodes_ .rank        ();

  // Make sure that the number of spatial dimensions (the rank of the
  // mesh) is valid.

  if ( rank_ < 1 || rank_ > 3 )
  {
    throw IllegalInputException (
      context,
      String::format (
        "invalid node rank: %d (should be 1, 2 or 3)", rank_
      )
    );
  }

  // Create an internal shape object for computing the element shape
  // functions.

  shape_ = IShapeFactory::newInstance (
    joinNames ( myName_, PropNames::SHAPE ),
    conf,
    props
  );  // create shape function related objects, with name myName_.SHAPE, looking for 
      // properties in props: myNam.shape

  // Make sure that the rank of the shape matches the rank of the
  // mesh.

  if ( shape_->globalRank() != rank_ )
  {
    throw IllegalInputException (
      context,
      String::format (
        "shape has invalid rank: %d (should be %d)",
        shape_->globalRank (),
        rank_
      )
    );
  }

  // Make sure that each element has the same number of nodes as the
  // shape object.

  elems_.checkSomeElements (
    context,
    egroup_.getIndices (),
    shape_->nodeCount  ()
  );

  // Get the degree of freedom (DOF) space, and add the DOFs of this
  // model.

  dofs_ = XDofSpace::get ( nodes_.getData(), globdat );  
                        // nodes_.getData() returns an object containing all the nodes
  dofTypes_.resize ( 1 );  //dofTypes_ contains column indices of dofs_ types
  //dofTypes_.resize ( rank_ );  //dofTypes_ contains column indices of dofs_ types
  // for ( int i = 0; i < rank_; i++ )
  // {
  //  dofTypes_[i] = dofs_->addType ( Names::DOFS[i] );
  // }

  dofTypes_ = dofs_->addType ( Names::DOFS[0] );
  dofs_->addDofs (
  elems_.getUniqueNodesOf ( egroup_.getIndices() ),
  dofTypes_
  ); // returns array containing indices of nodes attached to elements in egroup_

  // Create a material model object.

  material_      = newInstance<IsoMaterial> ( rank_ );  // choose material model
  getShapeGrads_ = getShapeGradsFunc        ( rank_ );  // create B matrix

  dummy_ = 1.;
  iterCount_ = 1.;
}

PhaseFieldModel::~PhaseFieldModel ()
{}


//-----------------------------------------------------------------------
//   configure
//-----------------------------------------------------------------------


void PhaseFieldModel::configure

  ( const Properties&  props,
    const Properties&  globdat )

{
  Properties  myProps = props.findProps ( myName_ );

  material_->configure ( myProps );
}


//-----------------------------------------------------------------------
//   getConfig
//-----------------------------------------------------------------------


void PhaseFieldModel::getConfig

  ( const Properties&  conf,
    const Properties&  globdat ) const

{
  Properties  myConf = conf.makeProps ( myName_ );

  material_->getConfig ( myConf );
}


//-----------------------------------------------------------------------
//   takeAction
//-----------------------------------------------------------------------


bool PhaseFieldModel::takeAction

  ( const String&      action,
    const Properties&  params,
    const Properties&  globdat )

{
  using jive::model::Actions;
  using jive::model::ActionParams;
  using jive::model::StateVector;


  if ( action == Actions::GET_MATRIX0 )
  {
    Ref<MatrixBuilder>  mbuilder;

    Vector  disp;
    Vector  force;

    // Get the current displacements.

    StateVector::get ( disp, dofs_, globdat );

    // Get the matrix builder and the internal force vector.

    params.get ( mbuilder, ActionParams::MATRIX0 );
    params.get ( force,    ActionParams::INT_VECTOR ); // pointing to the one originally created by the linsolve module. Initially set to zero

    getMatrix_ ( *mbuilder, force, disp );

    iterCount_ += 1;

    return true;
  }

  if ( action == Actions::GET_TABLE )
  {
    return getTable_ ( params, globdat );
  }



/*
  if ( action == Actions::COMMIT )
  {
    dummy_ = dummy_ + 0.4*dummy_;
    globdat.set ( "var.dummy", dummy_ );
  }


  if ( action == Actions::CHECK_CONVERGED )
  {
      Vector res;
      params.get ( res, ActionParams::RESIDUAL );
      jem::System::out() << "Norm of the residual = " << std::sqrt(dot(res,res)) << "\n";
      if (iterCount_ < 4)
      {
        params.set (ActionParams::CONVERGED, false);
      }
  }

  if ( action == Actions::ADVANCE )
  {
    iterCount_ = 0;
  }
*/

  return false;

}



//-----------------------------------------------------------------------
//   getMatrix_
//-----------------------------------------------------------------------


void PhaseFieldModel::getMatrix_

  ( MatrixBuilder&  mbuilder,
    const Vector&   force,
    const Vector&   disp ) const

{
  using jem::ALL;
  using jem::numeric::matmul;

  IntVector   ielems     = egroup_.getIndices  ();   // indices of elements in egroup_

  const int   ielemCount = ielems.size         ();
  const int   nodeCount  = shape_->nodeCount   ();
  const int   ipCount    = shape_->ipointCount ();   // integration point count
  const int   strCount   = STRAIN_COUNTS[rank_];
  const int   dofCount   = nodeCount;        // nb dof per element
  const double l = 0.5; // length-scale parameter

  Cubix       grads      ( rank_, nodeCount, ipCount );
  // 3D array of doubles, to store derivative of the shape function with respect to the global coordinates. grad(i,j,k) = i-th derivative of j-th shape function in k-th integration point
  
  Matrix      coords     ( rank_, nodeCount );

  Matrix      elemMat    ( dofCount, dofCount );
  Matrix      elemMat2    ( dofCount, dofCount );
  Vector      elemForce  ( dofCount );
  Vector      elemDisp   ( dofCount );

  Matrix      b          ( strCount, dofCount );
  Matrix      bt         = b.transpose ();

  Matrix      n          ( 1, 4 );
  Matrix      nt         = n.transpose ();

  IntVector   inodes     ( nodeCount );
  IntVector   idofs      ( dofCount  );

  Vector      ipWeights  ( ipCount   );

  MChain2     mc2;
  MChain3     mc3;

  
  // Iterate over all elements assigned to this model.

  for ( int ie = 0; ie < ielemCount; ie++ )
  {
    // Get the global element index.

    int  ielem = ielems[ie];

    // Get the element coordinates and DOFs.
    // functions in the integration points.

    elems_.getElemNodes  ( inodes, ielem    );

    nodes_.getSomeCoords ( coords, inodes );

    // dofTypes_[rank_-1] contains the dof for each node for the current dimension
    dofs_->getDofIndices ( idofs,  inodes, dofTypes_);

    
    // Compute the spatial derivatives of the element shape

    shape_->getShapeGradients ( grads, ipWeights, coords );

    Matrix shapes = shape_->getShapeFunctions ( );
    
    // Get the displacements at the element nodes.

    elemDisp = select ( disp, idofs );

    // Assemble the element matrix.

    elemMat = 0.0;

    for ( int ip = 0; ip < ipCount; ip++ )
    {
      // Compute the B-matrix for this integration point.

      getShapeGrads_ ( b, grads(ALL,ALL,ip) );
     
      // Compute N-matrix with shape functions for integration point 
      n(0,ALL) = shapes ( ALL,ip ) ;
      
      // Add the contribution of this integration point.

     elemMat +=  l * ipWeights[ip] *  mc2.matmul ( bt, b ) ;
     elemMat +=  1./l * ipWeights[ip]  * mc2.matmul ( nt, n );

      if (ie == 0)
      {
        //jem::System::out() << "!!!!!!!!!!!!!!!!!!elemMat2  = \n" << elemMat2 << "end\n" ;
      }

    }

    // Add the element matrix to the global stiffness matrix.

    mbuilder.addBlock ( idofs, idofs, elemMat );

    // Compute the element force vector.

    matmul ( elemForce, elemMat, elemDisp );

    // Add the element force vector to the global force vector.

    select ( force, idofs ) += elemForce;
  }
}


//-----------------------------------------------------------------------
//   getTable_
//-----------------------------------------------------------------------


bool PhaseFieldModel::getTable_

  ( const Properties&  params,
    const Properties&  globdat )

{
  using jive::model::Actions;
  using jive::model::ActionParams;
  using jive::model::StateVector;

  Ref<XTable>  table;
  Vector       weights;
  String       name;


  // Get the table, the name of the table, and the table row weights
  // from the action parameters.

  params.get ( table,   ActionParams::TABLE );
  params.get ( name,    ActionParams::TABLE_NAME );
  params.get ( weights, ActionParams::TABLE_WEIGHTS );

  // Stress values are computed in the nodes.

  if ( name == "stress" &&
       table->getRowItems() == nodes_.getData() )
  {
    Vector  disp;

    StateVector::get ( disp, dofs_, globdat );

    getStress_ ( *table, weights, disp );

    return true;
  }

  return false;
}


//-----------------------------------------------------------------------
//   getStress_
//-----------------------------------------------------------------------


void PhaseFieldModel::getStress_

  ( XTable&        table,
    const Vector&  weights,
    const Vector&  disp )

{
  using jem::ALL;
  using jem::numeric::matmul;

  IntVector  ielems     = egroup_.getIndices  ();

  const int  ielemCount = ielems.size         ();
  const int  nodeCount  = shape_->nodeCount   ();
  const int  strCount   = STRAIN_COUNTS[rank_];
  const int  dofCount   = rank_ * nodeCount;

  Cubix      grads      ( rank_, nodeCount, nodeCount );

  Matrix     stiff      = material_->getStiffMat ();

  Matrix     stress     ( nodeCount, strCount );
  Matrix     coords     ( rank_, nodeCount );
  Vector     elemDisp   ( dofCount );

  Matrix     b          ( strCount, dofCount );

  IntVector  inodes     ( nodeCount );
  IntVector  idofs      ( dofCount );
  IntVector  jcols      ( strCount );

  MChain2    mc2;


  // Add the columns for the stress components to the table.

  switch ( strCount )
  {
  case 1:

    jcols[0] = table.addColumn ( "sxx" );

    break;

  case 3:

    jcols[0] = table.addColumn ( "sxx" );
    jcols[1] = table.addColumn ( "syy" );
    jcols[2] = table.addColumn ( "sxy" );

    break;

  case 6:

    jcols[0] = table.addColumn ( "sxx" );
    jcols[1] = table.addColumn ( "syy" );
    jcols[2] = table.addColumn ( "szz" );
    jcols[3] = table.addColumn ( "sxy" );
    jcols[4] = table.addColumn ( "syz" );
    jcols[5] = table.addColumn ( "sxz" );

    break;

  default:

    throw jem::Error (
      JEM_FUNC,
      "unexpected number of stress components: " +
      String ( strCount )
    );
  }

  // Iterate over all elements assigned to this model.

  for ( int ie = 0; ie < ielemCount; ie++ )
  {
    // Get the global element index.

    int  ielem = ielems[ie];

    // Get the element coordinates and DOFs.

    elems_.getElemNodes  ( inodes, ielem  );
    nodes_.getSomeCoords ( coords, inodes );
    dofs_->getDofIndices ( idofs,  inodes, dofTypes_ );

    // Get the displacements in the nodes of this element.

    elemDisp = select ( disp, idofs );

    // Compute the spatial derivatives of the element shape
    // functions in the nodes.

    shape_->getVertexGradients ( grads, coords );

    // Compute the stresses in all nodes.

    for ( int i = 0; i < nodeCount; i++ )
    {
      // Compute the B-matrix for this node.

      getShapeGrads_ ( b, grads(ALL,ALL,i) );

      stress(i,ALL) = mc2.matmul ( stiff, b, elemDisp );
    }

    // Increment the table weights. When the complete table has been
    // filled, Jive will divide each row in the table by the
    // corresponding table weight. In this way the stress components
    // are automatically averaged over all elements that are attached
    // to a node. The weight vector is initially zero.

    select ( weights, inodes ) += 1.0;

    // Add the stresses to the table.

    table.addBlock ( inodes, jcols, stress );
  }
}


//=======================================================================
//   related functions
//=======================================================================

//-----------------------------------------------------------------------
//   newPhaseFieldModel
//-----------------------------------------------------------------------


static Ref<Model>     newPhaseFieldModel

  ( const String&       name,
    const Properties&   conf,
    const Properties&   props,
    const Properties&   globdat )

{
  //jem::System::out() << "%%%%%%%%%%%%\n" << props << "%%%%%%%%%%%%\n";
  return jem::newInstance<PhaseFieldModel> ( name,  conf,
                                          props, globdat );
}


//-----------------------------------------------------------------------
//   declarePhaseFieldModel
//-----------------------------------------------------------------------
// Register the new elastic model class newPhaseFieldModel

void declarePhaseFieldModel ()
{
  using jive::model::ModelFactory;

  ModelFactory::declare ( "Elastic", & newPhaseFieldModel );
  // i.e. use newPhaseFieldModel when jive wants to create an "Elastic" model, cf input file
}
